import numpy as np

from core.data.dataset import Dataset


class Spiral(Dataset):
    """
    螺旋数据集类，用于生成螺旋形状的数据。

    Attributes:
        data (numpy.ndarray): 输入数据，每个元素是螺旋形状上的点。
        label (numpy.ndarray): 目标标签，每个元素是对应点的类别。
    """

    def prepare(self):
        """
        准备数据集。生成螺旋形状的数据，包括训练和测试集。
        """
        self.data, self.label = self.get_spiral(self.train)

    def get_spiral(self, train=True):
        """
        生成螺旋形状的数据。

        Args:
            train (bool): 是否为训练集。默认为True。

        Returns:
            tuple: 包含输入数据和目标标签的元组。
        """
        # 设置随机种子
        seed = 1984 if train else 2020
        np.random.seed(seed=seed)

        # 定义数据集参数
        num_data, num_class, input_dim = 100, 3, 2
        data_size = num_class * num_data

        # 初始化数据和标签
        x = np.zeros((data_size, input_dim), dtype=np.float32)
        t = np.zeros(data_size, dtype=int)

        # 生成螺旋形状的数据
        for j in range(num_class):
            for i in range(num_data):
                rate = i / num_data
                radius = 1.0 * rate
                theta = j * 4.0 + 4.0 * rate + np.random.randn() * 0.2
                ix = num_data * j + i
                x[ix] = np.array([radius * np.sin(theta),
                                  radius * np.cos(theta)]).flatten()
                t[ix] = j

        # 打乱数据顺序
        indices = np.random.permutation(num_data * num_class)
        x = x[indices]
        t = t[indices]

        # 返回生成的数据和标签
        return x, t

